use ndarray::{Array1, Array2};

pub fn solve_by_direct(x: Array2<f64>, y: Array1<f64>) -> Array1<f64> {
    let n = x.nrows();
    assert_eq!(x.ncols(), n, "Matrix must be square");
    assert_eq!(y.len(), n, "Matrix and vector dimensions must match");

    let mut a = x;
    let mut b = y;

    // 前向消元
    for i in 0..n {
        // 选择主元并交换行
        let mut max_row = i;
        for k in (i + 1)..n {
            if a[[k, i]].abs() > a[[max_row, i]].abs() {
                max_row = k;
            }
        }
        if max_row != i {
            // 手动交换矩阵行
            for j in 0..n {
                a.swap((i, j), (max_row, j));
            }
            // 手动交换向量行
            b.swap(i, max_row);
        }

        // 消元
        for k in (i + 1)..n {
            let factor = a[[k, i]] / a[[i, i]];
            for j in i..n {
                a[[k, j]] -= factor * a[[i, j]];
            }
            b[k] -= factor * b[i];
        }
    }

    // 回代求解
    let mut x = Array1::<f64>::zeros(n);
    for i in (0..n).rev() {
        let mut sum = 0.0;
        for j in (i + 1)..n {
            sum += a[[i, j]] * x[j];
        }
        x[i] = (b[i] - sum) / a[[i, i]];
    }
    x
}

#[cfg(test)]
#[allow(unused)]
mod test {
    use super::*;
    use ndarray::array;
    use approx::assert_abs_diff_eq;

    #[test]
    fn test_solve() {
        let a = array![
            [4.0, 1.0, 2.0],
            [3.0, 5.0, 1.0],
            [1.0, 1.0, 3.0]
        ];
        let b = array![4.0, 7.0, 3.0];

        let x = solve_by_direct(a, b);
        
        let ans = array![0.5, 1.0, 0.5];
        
        assert_abs_diff_eq!(
            x.as_slice().unwrap(), 
            ans.as_slice().unwrap(),
            epsilon = 1e-9
        );
    }
}